from rbt.dmu import DecisionMakingUnit
import datetime


class KlineDMU(DecisionMakingUnit):
    """KlineDMU用于生成K线数据。
    不同交易所仅允许一个类型的KlineDMU实例，因此没有修改类名。
    """
    version = "v0"

    def __init__(self, interval=1):
        """
        初始化KlineDMU。

        参数:
            interval (int): K线的时间间隔，单位为分钟，默认为1分钟。
        """
        super().__init__()
        self.interval = interval  # K线时间间隔
        self.current_kline_end = None  # 当前K线开始时间
        self.kline_open = None  # K线开盘价
        self.kline_high = None  # K线最高价
        self.kline_low = None  # K线最低价
        self.kline_close = None  # K线收盘价
        self.kline_volume = None  # K线成交量
        self.kline_oi_diff = None  # K线持仓量变动
        self.md_count = 0  # 当前K线已接收的行情数
        self.update_end_time_required = True
        self.update_unit_name()  # 更新单元名称

    def get_param_str(self) -> str:
        """
        生成参数信息，用于更新单元名称。

        Returns:
            str: 载明参数的字符串。
        """
        return f"{self.interval}min"

    def make_decision(self, new_data, previous_result: dict = {}) -> dict:
        current_time = new_data.name  # 当前时间
        last_price = new_data["close"]  # 最终价
        cur_volume = new_data["volume"]  # 该行情戳的成交量

        # 仅在初始化第一根k线的结束时间时运行
        if self.current_kline_end is None:
            self.current_kline_end = datetime.datetime.combine(
                current_time.date(), datetime.time(0)
            )
            self.kline_open = last_price
            self.kline_high = last_price
            self.kline_low = last_price
            self.kline_close = last_price
            self.kline_volume = 0

        if self.update_end_time_required:
            while self.current_kline_end < current_time:
                self.current_kline_end += datetime.timedelta(minutes=self.interval)
            self.update_end_time_required = False

        # 如果k线走完且当前行情戳不计入k线
        res = None
        if self.current_kline_end < current_time:
            res = {
                "open": self.kline_open,
                "high": self.kline_high,
                "low": self.kline_low,
                "close": self.kline_close,
                "volume": self.kline_volume,
                "end_time": self.current_kline_end,
                "md_count": self.md_count,
                "completed": True,
            }
            while self.current_kline_end < current_time:
                self.current_kline_end += datetime.timedelta(minutes=self.interval)
            self.kline_open = last_price
            self.kline_high = last_price
            self.kline_low = last_price
            self.kline_close = last_price
            self.kline_volume = cur_volume
            return res

        # 如果当前行情戳计入K线
        self.kline_open = last_price if self.kline_open is None else self.kline_open
        self.kline_close = last_price
        self.kline_high = (
            last_price if self.kline_high is None else max(self.kline_high, last_price)
        )
        self.kline_low = (
            last_price if self.kline_low is None else min(self.kline_low, last_price)
        )
        self.kline_volume += cur_volume
        self.md_count += 1

        if self.current_kline_end == current_time:
            res = {
                "open": self.kline_open,
                "high": self.kline_high,
                "low": self.kline_low,
                "close": self.kline_close,
                "volume": self.kline_volume,
                "end_time": self.current_kline_end,
                "md_count": self.md_count,
                "completed": True,
            }
            self.kline_open = None
            self.kline_close = None
            self.kline_high = None
            self.kline_low = None
            self.update_end_time_required = True
            self.md_count = 0
            return res

        res = {
            "open": self.kline_open,
            "high": self.kline_high,
            "low": self.kline_low,
            "close": self.kline_close,
            "volume": self.kline_volume,
            "end_time": self.current_kline_end,
            "md_count": self.md_count,
            "completed": False,
        }
        return res
